#include <stdio.h>
#include <stdlib.h>
#include <omp.h>
#include <time.h>

// 初始化矩阵
void init_matrix(double ***matrix, int rows, int cols) {
    *matrix = (double **)malloc(rows * sizeof(double *));
    for (int i = 0; i < rows; i++) {
        (*matrix)[i] = (double *)malloc(cols * sizeof(double));
        for (int j = 0; j < cols; j++) {
            (*matrix)[i][j] = (double)rand() / RAND_MAX * 100.0;
        }
    }
}

// 打印矩阵（前10x10部分）
void print_matrix(double **matrix, int rows, int cols) {
    int print_rows = rows > 10 ? 10 : rows;
    int print_cols = cols > 10 ? 10 : cols;
    
    for (int i = 0; i < print_rows; i++) {
        for (int j = 0; j < print_cols; j++) {
            printf("%8.2f ", matrix[i][j]);
        }
        printf("\n");
    }
}

// 释放矩阵内存
void free_matrix(double **matrix, int rows) {
    for (int i = 0; i < rows; i++) {
        free(matrix[i]);
    }
    free(matrix);
}

int main(int argc, char *argv[]) {
    if (argc < 5 || argc > 6) {
        printf("Usage: %s <m> <n> <k> <num_threads> [schedule]\n", argv[0]);
        printf("Schedule options: 0-default, 1-static, 2-dynamic\n");
        return 1;
    }
    
    int m = atoi(argv[1]);
    int n = atoi(argv[2]);
    int k = atoi(argv[3]);
    int num_threads = atoi(argv[4]);
    int schedule = 0; // 默认调度
    
    if (argc > 5) {
        schedule = atoi(argv[5]);
    }
    
    if (m < 128 || m > 2048 || n < 128 || n > 2048 || k < 128 || k > 2048) {
        printf("Matrix dimensions must be between 128 and 2048\n");
        return 1;
    }
    
    if (num_threads < 1 || num_threads > 16) {
        printf("Number of threads must be between 1 and 16\n");
        return 1;
    }
    
    srand(time(NULL));
    
    // 初始化矩阵
    double **A, **B, **C;
    init_matrix(&A, m, n);
    init_matrix(&B, n, k);
    init_matrix(&C, m, k);
    
    // 设置线程数
    omp_set_num_threads(num_threads);
    
    double start_time = omp_get_wtime();

    // 根据调度模式设置并行计算
switch(schedule) {
    case 0: // 默认调度
        #pragma omp parallel for
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < k; j++) {
                C[i][j] = 0.0;
                for (int l = 0; l < n; l++) {
                    C[i][j] += A[i][l] * B[l][j];
                }
            }
        }
        break;
        
    case 1: // 静态调度
        #pragma omp parallel for schedule(static)
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < k; j++) {
                C[i][j] = 0.0;
                for (int l = 0; l < n; l++) {
                    C[i][j] += A[i][l] * B[l][j];
                }
            }
        }
        break;
        
    case 2: // 动态调度
        #pragma omp parallel for schedule(dynamic)
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < k; j++) {
                C[i][j] = 0.0;
                for (int l = 0; l < n; l++) {
                    C[i][j] += A[i][l] * B[l][j];
                }
            }
        }
        break;
}

    double end_time = omp_get_wtime();
    double elapsed_time = end_time - start_time;
    
    // 输出结果
    printf("Matrix A (first 10x10):\n");
    print_matrix(A, m, n);
    
    printf("\nMatrix B (first 10x10):\n");
    print_matrix(B, n, k);
    
    printf("\nMatrix C (Result, first 10x10):\n");
    print_matrix(C, m, k);
    
    printf("\nTime taken: %.6f seconds\n", elapsed_time);
    printf("Matrix size: %dx%dx%d, Threads: %d, Schedule: %d\n", 
           m, n, k, num_threads, schedule);
    
    // 释放内存
    free_matrix(A, m);
    free_matrix(B, n);
    free_matrix(C, m);
    
    return 0;
}